#!/usr/bin/python
from freenect import sync_get_depth as get_depth, sync_get_video as get_video
import math
import cv  
import numpy as np
import Image
import networkx as nx
import random
from time import sleep

#Some constants 
WAIT_TIME = 5 # CV time wait in milliseconds
FORCE_MODIFIER = .0000001
TEN = .01
FIVE_SQRT = TEN*.5
NUDGE = .01*TEN
OMEGA_PARAM = 2
PHI_P_PARAM = 3
PHI_G_PARAM = 5
KINECT_WIDTH = 640
KINECT_HEIGHT = 480

totalRow=KINECT_HEIGHT
totalCol=KINECT_WIDTH
totalDepth=2047
boxSize=60
l = .4
lastcasenum= 0 
minDistance = -10
scaleFactor = .0021

def nudge():
    return random.uniform(-NUDGE,NUDGE)

def D2M(z):
        return (1000.0 / (z * -0.0030711016 + 3.3309495161));

class Particle:
    def move(self):
        if self.moves:
            self.setCoords(self.x + (self.velx)*WAIT_TIME,
                           self.y - (self.vely)*WAIT_TIME)
    
    def depth(self,newX,newY):
        #print newX, newY
        (depth,_) = get_depth()
        #print "-----"
        return D2M(depth[newY][newX])
    
    def setCoords(self,newX,newY):
        #print newX, newY
    	self._x = math.floor(newX%KINECT_WIDTH)
    	self._y = math.floor(newY%KINECT_HEIGHT)
    	self._z = self.depth(self._x,self._y)
    
    def get_x(self):
        return self._x
    def get_y(self):
        return self._y
    def get_z(self):
        return self._z
    x = property(get_x)
    y = property(get_y)
    z = property(get_z)
    
    def __init__(self,xinit,yinit,ident):
    	self.velx = 0
    	self.vely = 0
        self.setCoords(xinit,yinit)
        
        self.moves = False
        self.i = ident
        
        self.bestCost = None
    
    def bestFit(self):
        return self
    
    def dist_sq(self,anotherParticle):
        return (self.x-anotherParticle.x)**2 + (self.y-anotherParticle.y)**2+ (self.z-anotherParticle.z)**2

class MovingParticle(Particle):
    def __init__(self,ident):
        self.velx = random.randint(-KINECT_WIDTH,KINECT_WIDTH)
        self.vely = random.randint(-KINECT_HEIGHT,KINECT_HEIGHT)
        self._x = random.randint(0,KINECT_WIDTH)
        self._y = random.randint(0,KINECT_HEIGHT)
        self.setCoords(self._x,self._y)
        self.p = Particle(self._x,self._y,-1)
        
        self.moves = True
        self.i = ident
        
        self.bestCost = None
    
    def bestFit(self):
    	return self.p

def pretty_depth(depth):
    """Converts depth into a 'nicer' format for display

    This is abstracted to allow for experimentation with normalization

    Args:
        depth: A numpy array with 2 bytes per pixel

    Returns:
        A numpy array that has been processed whos datatype is unspecified
    """
    np.clip(depth, 0, 2**10 - 1, depth)
    depth >>= 2
    depth = depth.astype(np.uint8)
    return depth
def getRowCol(flat):
      maxRow=math.floor(flat/KINECT_WIDTH)
      maxCol=flat % KINECT_WIDTH
      return [maxRow,maxCol]
def toFlat(row,col):
      return int((row%KINECT_HEIGHT)  * KINECT_WIDTH + (col%KINECT_WIDTH)) 
def colorPix(r,g,b,colorIndex):
      (theRow,theCol)=getRowCol(colorIndex)
      for n in range(0, 10):
         for m in range(0, 10):
        
                d3[(theRow+n)%KINECT_HEIGHT][(theCol+m)%KINECT_WIDTH][0]=r
                d3[(theRow+n)%KINECT_HEIGHT][(theCol+m)%KINECT_WIDTH][1]=g
                d3[(theRow+n)%KINECT_HEIGHT][(theCol+m)%KINECT_WIDTH][2]=b

def cost(node,model):
    errorSumX=0
    errorSumY=0
    errorSumZ=0
    for neighbor in model.neighbors(node):
        preferredDistance = model[node][neighbor]['dist']
        actualDistance = node.dist_sq(neighbor)
        xdiff = node.x - neighbor.x
        ydiff = node.y - neighbor.y
        zdiff = node.z - neighbor.z
        errorSumX += math.fabs(preferredDistance-actualDistance)*(xdiff**2)
        errorSumY += math.fabs(preferredDistance-actualDistance)*(ydiff**2)
        errorSumZ += math.fabs(preferredDistance-actualDistance)*(zdiff**2)
    return errorSumX**2 + errorSumY**2 + errorSumZ**2

def main():
       global d3, depth, rg ,ib,minIndex,maxIndex,min,nRow,nCol,nIndex,maxRow,maxCol
       
       #The process of model creation may be automated... later.
       #For now, it shall be explicit.
       model = nx.Graph()
       print "Five seconds. Get into position."
       sleep(5)
       (depth,_) = get_depth()
       #depth= pretty_depth(depth) 
       fDepth=depth.flatten()
       maxIndex= np.argmin(fDepth)
       
       (initMaxRow, initMaxCol) = getRowCol(maxIndex)
       centerParticle = Particle(initMaxRow,initMaxCol,-1)
       model.add_node(centerParticle)
       corners = []
       sides = []
       
       # Edge values represent distance squared.
       for i in range(6):
            if i < 3:
                #add corners
                corners.append(MovingParticle(i))
                model.add_node(corners[i])
                # each corner should be 10 units away from the center
                model.add_edge(corners[i],centerParticle,dist=TEN) 
            else:
                sides.append(MovingParticle(i))
                model.add_node(sides[i-3])
                #Side nodes are 5*sqrt(2) away from a corner and the center
                # 5*sqrt(2) is about 7.07106781
                model.add_edge(sides[i-3],centerParticle,dist=FIVE_SQRT)
                model.add_edge(sides[i-3],corners[i-3],dist=FIVE_SQRT)
                model.add_edge(sides[i-3],corners[(i-2)%3],dist=FIVE_SQRT)
       
       #End model creation.
       print "Looping time."
       #MAIN LOOP
       while True:
        # Get a fresh frame
        (depth,_) = get_depth()
        #depth= pretty_depth(depth) 
        fDepth=depth.flatten()
        maxIndex= np.argmin(fDepth)
        (initMaxRow, initMaxCol) = getRowCol(maxIndex)
        centerParticle.setCoords(initMaxRow, initMaxCol) #!!! Does this affect the particle in the graph?        
        
        d3 = np.dstack((depth,depth,depth)).astype(np.uint8)
        # Build a two panel color image
        colorPix(255,0,0,maxIndex)
        
        nodeDisplay = []
        
        for node in model.nodes():
            rp = random.random()
            node.vx = OMEGA_PARAM * node.velx + PHI_P_PARAM * rp * (node.bestFit().x - node.x)
            node.vy = OMEGA_PARAM * node.vely + PHI_P_PARAM * rp * (node.bestFit().y - node.y)
            node.move()
            nodeCost = cost(node,model)
            if node.bestCost == None or nodeCost < node.bestCost:
                node.p = Particle(node.x, node.y,-1)
                node.bestCost = nodeCost
            nodeDisplay.append(toFlat(node.x,node.y))
        
        for pix in nodeDisplay:
            colorPix(0,0,255,pix)
        
        
        cv.ShowImage('both',cv.fromarray(np.array(d3[::2,::2,::-1])))
        cv.WaitKey(WAIT_TIME)

         
main()

"""
IPython usage:
 ipython
 [1]: run -i demo_freenect
 #<ctrl -c>  (to interrupt the loop)
 [2]: %timeit -n100 get_depth(), get_rgb() # profile the kinect capture

"""

